use crate::ast::*;
use crate::error::{CompilerError, Result};
use crate::llvm_backend::LLVMBackend;
use inkwell::values::BasicValueEnum;

impl<'ctx> LLVMBackend<'ctx> {


    /// Generate map() method - applies function to Ok value, keeps Err
    pub(super) fn generate_result_map(
        &mut self,
        result_value: BasicValueEnum<'ctx>,
        result_type: &crate::ast::types::ResultType,
        arguments: &[Expression],
        location: &crate::error::Location,
    ) -> Result<BasicValueEnum<'ctx>> {
        if arguments.len() != 1 {
            return Err(CompilerError::type_error(
                location.line,
                location.column,
                &format!(
                    "Result.map() expects 1 argument (function), got {}",
                    arguments.len()
                ),
            ));
        }

        let func = self
            .current_function
            .ok_or_else(|| CompilerError::internal("Cannot create blocks outside function"))?;

        let tag = self.extract_result_tag(result_value)?;
        let zero = self.context.i32_type().const_int(0, false);
        let is_ok =
            self.builder
                .build_int_compare(inkwell::IntPredicate::EQ, tag, zero, "map_check")?;

        // Create blocks: ok_block (apply function), err_block (keep Err)
        let ok_block = self.context.append_basic_block(func, "map_ok");
        let err_block = self.context.append_basic_block(func, "map_err");
        let continue_block = self.context.append_basic_block(func, "map_continue");

        self.builder
            .build_conditional_branch(is_ok, ok_block, err_block)?;

        // Ok branch: extract Ok value, call function, create new Result::Ok
        self.builder.position_at_end(ok_block);
        let ok_value = self.extract_result_ok_value(result_value, result_type)?;

        // Store ok_value in a temporary variable so we can use it in function call
        let ok_type = self.nrc_type_to_llvm_type(&Some(*result_type.ok_type.clone()))?;
        let ok_temp_alloca = self.builder.build_alloca(ok_type, "ok_temp")?;
        self.builder.build_store(ok_temp_alloca, ok_value)?;

        // Create a temporary variable name for the ok value
        let ok_temp_var = format!("__map_ok_{}", location.line);

        // Store the alloca in variable_map temporarily
        let old_ok_value = self
            .variable_map
            .insert(ok_temp_var.clone(), ok_temp_alloca.into());
        let old_ok_type = self
            .variable_type_map
            .insert(ok_temp_var.clone(), *result_type.ok_type.clone());

        // Create a call expression: f(ok_temp_var)
        let func_expr = arguments[0].clone();
        let ok_var_expr = Expression::Variable(ok_temp_var.clone());
        let call_expr = Expression::Call(CallExpr {
            callee: Box::new(func_expr),
            arguments: vec![ok_var_expr],
            type_args: Vec::new(),
            resolved: None,
            location: *location,
        });

        // Generate the function call
        let transformed_value = self.generate_expression(&call_expr)?;

        // Infer the transformed value's type from the function return type
        // For now, we'll use the function's return type from type inference
        // TODO: Get the actual return type from type inference
        let transformed_type = self
            .infer_expression_type(&call_expr)
            .unwrap_or_else(|| *result_type.ok_type.clone()); // Fallback to Ok type

        // Restore variable map
        if let Some(old) = old_ok_value {
            self.variable_map.insert(ok_temp_var.clone(), old);
        } else {
            self.variable_map.remove(&ok_temp_var);
        }
        if let Some(old) = old_ok_type {
            self.variable_type_map.insert(ok_temp_var, old);
        } else {
            self.variable_type_map.remove(&ok_temp_var);
        }

        // Create new Result::Ok with transformed value
        // The new Result type should be Result<U, E> where U is the function's return type
        // For now, we'll use the same result_type structure but with transformed ok_type
        let new_result_type = crate::ast::types::ResultType {
            ok_type: Box::new(transformed_type.clone()),
            err_type: result_type.err_type.clone(),
            location: *location,
        };
        let new_ok_result =
            self.create_result_value(transformed_value, &transformed_type, &new_result_type, "Ok")?;

        self.builder.build_unconditional_branch(continue_block)?;

        // Err branch: create new Result::Err with same error value
        self.builder.position_at_end(err_block);
        let err_value = self.extract_result_err_value(result_value, result_type)?;
        let new_err_result =
            self.create_result_value(err_value, &*result_type.err_type, result_type, "Err")?;
        self.builder.build_unconditional_branch(continue_block)?;

        // Continue block: return either new Ok Result or new Err Result
        self.builder.position_at_end(continue_block);
        let result_llvm_type = self.get_result_llvm_type(&new_result_type)?;
        let phi = self.builder.build_phi(result_llvm_type, "map_result")?;
        phi.add_incoming(&[(&new_ok_result, ok_block), (&new_err_result, err_block)]);
        Ok(phi.as_basic_value())
    }


    /// Generate map_err() method - applies function to Err value, keeps Ok
    pub(super) fn generate_result_map_err(
        &mut self,
        result_value: BasicValueEnum<'ctx>,
        result_type: &crate::ast::types::ResultType,
        arguments: &[Expression],
        location: &crate::error::Location,
    ) -> Result<BasicValueEnum<'ctx>> {
        if arguments.len() != 1 {
            return Err(CompilerError::type_error(
                location.line,
                location.column,
                &format!(
                    "Result.map_err() expects 1 argument (function), got {}",
                    arguments.len()
                ),
            ));
        }

        let func = self
            .current_function
            .ok_or_else(|| CompilerError::internal("Cannot create blocks outside function"))?;

        let tag = self.extract_result_tag(result_value)?;
        let zero = self.context.i32_type().const_int(0, false);
        let is_ok = self.builder.build_int_compare(
            inkwell::IntPredicate::EQ,
            tag,
            zero,
            "map_err_check",
        )?;

        // Create blocks: ok_block (keep Ok), err_block (apply function)
        let ok_block = self.context.append_basic_block(func, "map_err_ok");
        let err_block = self.context.append_basic_block(func, "map_err_err");
        let continue_block = self.context.append_basic_block(func, "map_err_continue");

        self.builder
            .build_conditional_branch(is_ok, ok_block, err_block)?;

        // Ok branch: create new Result::Ok with same Ok value
        self.builder.position_at_end(ok_block);
        let ok_value = self.extract_result_ok_value(result_value, result_type)?;
        let new_ok_result =
            self.create_result_value(ok_value, &*result_type.ok_type, result_type, "Ok")?;
        self.builder.build_unconditional_branch(continue_block)?;

        // Err branch: extract Err value, call function, create new Result::Err
        self.builder.position_at_end(err_block);
        let err_value = self.extract_result_err_value(result_value, result_type)?;

        // Store err_value in a temporary variable so we can use it in function call
        let err_type = self.nrc_type_to_llvm_type(&Some(*result_type.err_type.clone()))?;
        let err_temp_alloca = self.builder.build_alloca(err_type, "err_temp")?;
        self.builder.build_store(err_temp_alloca, err_value)?;

        // Create a temporary variable name for the err value
        let err_temp_var = format!("__map_err_err_{}", location.line);

        // Store the alloca in variable_map temporarily
        let old_err_value = self
            .variable_map
            .insert(err_temp_var.clone(), err_temp_alloca.into());
        let old_err_type = self
            .variable_type_map
            .insert(err_temp_var.clone(), *result_type.err_type.clone());

        // Create a call expression: f(err_temp_var)
        let func_expr = arguments[0].clone();
        let err_var_expr = Expression::Variable(err_temp_var.clone());
        let call_expr = Expression::Call(CallExpr {
            callee: Box::new(func_expr),
            arguments: vec![err_var_expr],
            type_args: Vec::new(),
            resolved: None,
            location: *location,
        });

        // Generate the function call
        let transformed_value = self.generate_expression(&call_expr)?;

        // Infer the transformed value's type from the function return type
        let transformed_type = self
            .infer_expression_type(&call_expr)
            .unwrap_or_else(|| *result_type.err_type.clone()); // Fallback to Err type

        // Restore variable map
        if let Some(old) = old_err_value {
            self.variable_map.insert(err_temp_var.clone(), old);
        } else {
            self.variable_map.remove(&err_temp_var);
        }
        if let Some(old) = old_err_type {
            self.variable_type_map.insert(err_temp_var, old);
        } else {
            self.variable_type_map.remove(&err_temp_var);
        }

        // Create new Result::Err with transformed value
        // The new Result type should be Result<T, F> where F is the function's return type
        let new_result_type = crate::ast::types::ResultType {
            ok_type: result_type.ok_type.clone(),
            err_type: Box::new(transformed_type.clone()),
            location: *location,
        };
        let new_err_result = self.create_result_value(
            transformed_value,
            &transformed_type,
            &new_result_type,
            "Err",
        )?;

        self.builder.build_unconditional_branch(continue_block)?;

        // Continue block: return either new Ok Result or new Err Result
        self.builder.position_at_end(continue_block);
        let result_llvm_type = self.get_result_llvm_type(&new_result_type)?;
        let phi = self.builder.build_phi(result_llvm_type, "map_err_result")?;
        phi.add_incoming(&[(&new_ok_result, ok_block), (&new_err_result, err_block)]);
        Ok(phi.as_basic_value())
    }


    /// Generate map_or() method - maps Ok or returns default
    /// Ok -> f(value), Err -> default
    pub(super) fn generate_result_map_or(
        &mut self,
        result_value: BasicValueEnum<'ctx>,
        result_type: &crate::ast::types::ResultType,
        arguments: &[Expression],
        location: &crate::error::Location,
    ) -> Result<BasicValueEnum<'ctx>> {
        if arguments.len() != 2 {
            return Err(CompilerError::type_error(
                location.line,
                location.column,
                &format!(
                    "Result.map_or() expects 2 arguments (default, function), got {}",
                    arguments.len()
                ),
            ));
        }

        let func = self
            .current_function
            .ok_or_else(|| CompilerError::internal("Cannot create blocks outside function"))?;

        // Generate default value
        let default_value = self.generate_expression(&arguments[0])?;
        let default_type = match self.infer_expression_type(&arguments[0]) {
            Some(t) => t,
            _ => return Err(CompilerError::internal("Could not infer default type")),
        };

        let tag = self.extract_result_tag(result_value)?;
        let zero = self.context.i32_type().const_int(0, false);
        let is_ok =
            self.builder
                .build_int_compare(inkwell::IntPredicate::EQ, tag, zero, "map_or_check")?;

        // Create blocks: ok_block (call function), err_block (return default)
        let ok_block = self.context.append_basic_block(func, "map_or_ok");
        let err_block = self.context.append_basic_block(func, "map_or_err");
        let continue_block = self.context.append_basic_block(func, "map_or_continue");

        self.builder
            .build_conditional_branch(is_ok, ok_block, err_block)?;

        // Ok block: extract value, call function
        self.builder.position_at_end(ok_block);
        let ok_value = self.extract_result_ok_value(result_value, result_type)?;

        let ok_type = self.nrc_type_to_llvm_type(&Some(*result_type.ok_type.clone()))?;
        let ok_temp_alloca = self.builder.build_alloca(ok_type, "ok_temp")?;
        self.builder.build_store(ok_temp_alloca, ok_value)?;

        let ok_temp_var = format!("__map_or_ok_{}", location.line);
        let old_ok_value = self
            .variable_map
            .insert(ok_temp_var.clone(), ok_temp_alloca.into());
        let old_ok_type = self
            .variable_type_map
            .insert(ok_temp_var.clone(), *result_type.ok_type.clone());

        // Call function
        let func_expr = arguments[1].clone();
        let ok_var_expr = Expression::Variable(ok_temp_var.clone());
        let call_expr = Expression::Call(CallExpr {
            callee: Box::new(func_expr),
            arguments: vec![ok_var_expr],
            type_args: Vec::new(),
            resolved: None,
            location: *location,
        });

        let mapped_value = self.generate_expression(&call_expr)?;

        // Restore variable map
        if let Some(old) = old_ok_value {
            self.variable_map.insert(ok_temp_var.clone(), old);
        } else {
            self.variable_map.remove(&ok_temp_var);
        }
        if let Some(old) = old_ok_type {
            self.variable_type_map.insert(ok_temp_var, old);
        } else {
            self.variable_type_map.remove(&ok_temp_var);
        }

        self.builder.build_unconditional_branch(continue_block)?;

        // Err block: return default
        self.builder.position_at_end(err_block);
        self.builder.build_unconditional_branch(continue_block)?;

        // Continue block: return either mapped value or default
        self.builder.position_at_end(continue_block);
        let result_llvm_type = self.nrc_type_to_llvm_type(&Some(default_type.clone()))?;
        let phi = self.builder.build_phi(result_llvm_type, "map_or_result")?;
        phi.add_incoming(&[(&mapped_value, ok_block), (&default_value, err_block)]);
        Ok(phi.as_basic_value())
    }


    /// Generate map_or_else() method - maps Ok or computes default
    /// Ok -> f(value), Err -> default_fn(error)
    pub(super) fn generate_result_map_or_else(
        &mut self,
        result_value: BasicValueEnum<'ctx>,
        result_type: &crate::ast::types::ResultType,
        arguments: &[Expression],
        location: &crate::error::Location,
    ) -> Result<BasicValueEnum<'ctx>> {
        if arguments.len() != 2 {
            return Err(CompilerError::type_error(
                location.line,
                location.column,
                &format!(
                    "Result.map_or_else() expects 2 arguments (default_fn, function), got {}",
                    arguments.len()
                ),
            ));
        }

        let func = self
            .current_function
            .ok_or_else(|| CompilerError::internal("Cannot create blocks outside function"))?;

        let tag = self.extract_result_tag(result_value)?;
        let zero = self.context.i32_type().const_int(0, false);
        let is_ok = self.builder.build_int_compare(
            inkwell::IntPredicate::EQ,
            tag,
            zero,
            "map_or_else_check",
        )?;

        // Create blocks: ok_block (call function), err_block (call default_fn)
        let ok_block = self.context.append_basic_block(func, "map_or_else_ok");
        let err_block = self.context.append_basic_block(func, "map_or_else_err");
        let continue_block = self
            .context
            .append_basic_block(func, "map_or_else_continue");

        self.builder
            .build_conditional_branch(is_ok, ok_block, err_block)?;

        // Ok block: extract value, call function
        self.builder.position_at_end(ok_block);
        let ok_value = self.extract_result_ok_value(result_value, result_type)?;

        let ok_type = self.nrc_type_to_llvm_type(&Some(*result_type.ok_type.clone()))?;
        let ok_temp_alloca = self.builder.build_alloca(ok_type, "ok_temp")?;
        self.builder.build_store(ok_temp_alloca, ok_value)?;

        let ok_temp_var = format!("__map_or_else_ok_{}", location.line);
        let old_ok_value = self
            .variable_map
            .insert(ok_temp_var.clone(), ok_temp_alloca.into());
        let old_ok_type = self
            .variable_type_map
            .insert(ok_temp_var.clone(), *result_type.ok_type.clone());

        // Call function
        let func_expr = arguments[1].clone();
        let ok_var_expr = Expression::Variable(ok_temp_var.clone());
        let call_expr = Expression::Call(CallExpr {
            callee: Box::new(func_expr),
            arguments: vec![ok_var_expr],
            type_args: Vec::new(),
            resolved: None,
            location: *location,
        });

        let mapped_value = self.generate_expression(&call_expr)?;

        // Restore variable map
        if let Some(old) = old_ok_value {
            self.variable_map.insert(ok_temp_var.clone(), old);
        } else {
            self.variable_map.remove(&ok_temp_var);
        }
        if let Some(old) = old_ok_type {
            self.variable_type_map.insert(ok_temp_var, old);
        } else {
            self.variable_type_map.remove(&ok_temp_var);
        }

        self.builder.build_unconditional_branch(continue_block)?;

        // Err block: extract error value, call default_fn
        self.builder.position_at_end(err_block);
        let err_value = self.extract_result_err_value(result_value, result_type)?;

        let err_type = self.nrc_type_to_llvm_type(&Some(*result_type.err_type.clone()))?;
        let err_temp_alloca = self.builder.build_alloca(err_type, "err_temp")?;
        self.builder.build_store(err_temp_alloca, err_value)?;

        let err_temp_var = format!("__map_or_else_err_{}", location.line);
        let old_err_value = self
            .variable_map
            .insert(err_temp_var.clone(), err_temp_alloca.into());
        let old_err_type = self
            .variable_type_map
            .insert(err_temp_var.clone(), *result_type.err_type.clone());

        // Call default_fn with error value
        let default_fn_expr = arguments[0].clone();
        let err_var_expr = Expression::Variable(err_temp_var.clone());
        let default_call_expr = Expression::Call(CallExpr {
            callee: Box::new(default_fn_expr),
            arguments: vec![err_var_expr],
            type_args: Vec::new(),
            resolved: None,
            location: *location,
        });

        let default_value = self.generate_expression(&default_call_expr)?;

        // Restore variable map
        if let Some(old) = old_err_value {
            self.variable_map.insert(err_temp_var.clone(), old);
        } else {
            self.variable_map.remove(&err_temp_var);
        }
        if let Some(old) = old_err_type {
            self.variable_type_map.insert(err_temp_var, old);
        } else {
            self.variable_type_map.remove(&err_temp_var);
        }

        // Infer return type from the function call (should match default_fn return type)
        let return_type = match self.infer_expression_type(&call_expr) {
            Some(t) => t,
            _ => match self.infer_expression_type(&default_call_expr) {
                Some(t) => t,
                _ => return Err(CompilerError::internal("Could not infer return type")),
            },
        };

        self.builder.build_unconditional_branch(continue_block)?;

        // Continue block: return either mapped value or default value
        self.builder.position_at_end(continue_block);
        let result_llvm_type = self.nrc_type_to_llvm_type(&Some(return_type.clone()))?;
        let phi = self
            .builder
            .build_phi(result_llvm_type, "map_or_else_result")?;
        phi.add_incoming(&[(&mapped_value, ok_block), (&default_value, err_block)]);
        Ok(phi.as_basic_value())
    }

}
